#!/usr/bin/env python3
"""
LGBM客户购买预测脚本 - 正确用法
"""
import pandas as pd
import numpy as np
import lightgbm as lgb
from sklearn.model_selection import KFold
from sklearn.metrics import roc_auc_score

print("=== LGBM客户购买预测 ===")

# 读取数据
print("读取数据...")
train = pd.read_csv('train_set.csv')
test = pd.read_csv('test_set.csv')

# 特征工程
print("特征工程...")
X = train.drop(['ID', 'y'], axis=1)
y = train['y']
X_test = test.drop(['ID'], axis=1)

# 对分类变量进行编码
categorical_features = ['job', 'marital', 'education', 'default', 'housing', 'loan', 'contact', 'month', 'poutcome']
for col in categorical_features:
    X[col] = X[col].astype('category')
    X_test[col] = X_test[col].astype('category')

# 5折交叉验证
print("开始LGBM模型训练...")
n_folds = 5
kf = KFold(n_splits=n_folds, shuffle=True, random_state=42)

oof_preds = np.zeros(len(X))
test_preds = np.zeros(len(X_test))

params = {
    'objective': 'binary',
    'metric': 'auc',
    'boosting_type': 'gbdt',
    'learning_rate': 0.05,
    'num_leaves': 31,
    'max_depth': -1,
    'min_child_samples': 20,
    'reg_alpha': 0.0,
    'reg_lambda': 0.0,
    'random_state': 42,
    'n_jobs': -1,
    'verbose': -1
}

for fold, (train_idx, val_idx) in enumerate(kf.split(X)):
    print(f"训练第 {fold+1}/{n_folds} 折")
    
    X_train, X_val = X.iloc[train_idx], X.iloc[val_idx]
    y_train, y_val = y.iloc[train_idx], y.iloc[val_idx]
    
    # 创建LGBM数据集
    train_data = lgb.Dataset(X_train, label=y_train, categorical_feature=categorical_features)
    val_data = lgb.Dataset(X_val, label=y_val, categorical_feature=categorical_features)
    
    # 训练模型
    model = lgb.train(
        params,
        train_data,
        valid_sets=[train_data, val_data],
        callbacks=[
            lgb.log_evaluation(50),
            lgb.early_stopping(200)
        ]
    )
    
    # 预测
    oof_preds[val_idx] = model.predict(X_val)
    test_preds += model.predict(X_test) / n_folds

# 计算AUC分数
auc_score = roc_auc_score(y, oof_preds)
print(f"✅ LGBM模型训练完成!")
print(f"📊 AUC 分数: {auc_score:.6f}")

# 生成预测文件
submission = pd.DataFrame({
    'ID': test['ID'],
    'pred': test_preds
})

output_file = f'lgbm_submission_auc_{auc_score:.6f}.csv'
submission.to_csv(output_file, index=False)
print(f"📁 预测文件已保存: {output_file}")
print("🎉 任务完成!")